import random

import numpy as np
import pandas as pd
from boruta import BorutaPy
from sklearn.ensemble import RandomForestRegressor
def fs_by_boruta(model, x, y, features=None, verbose=2, random_state=None):
    """
    feture selection by boruta
    x: np array
    y: np array
    self.selected_features to get features
    """
    if features is None:
        features = np.array([f"f{i}" for i in range(X.shape[1])])
    # define Boruta feature selection method
    feat_selector = BorutaPy(model, n_estimators='auto', verbose=verbose, random_state=random_state)
    feat_selector.fit(x, y)
    #X_filtered = feat_selector.transform(X)
    selected_features = features[feat_selector.support_]
    feat_selector.selected_features = selected_features
    return feat_selector


if __name__ == '__main__':
    random.seed(0)
    np.random.seed(0)
    Y_col = 'temperature'
    dataset = pd.read_csv(f"../data/ml_dataset_{Y_col}.csv")

    features = list(dataset.columns[:-1])
    X = dataset[features]
    Y = dataset[dataset.columns[-1]]
    model = RandomForestRegressor()
    fselect = fs_by_boruta(model, X.values, Y.values, features=features)
